# -*- coding: UTF-8 -*-

# %% [markdown]
# 
import sympy as sp
import numpy as np
import matplotlib.pyplot as plt
import scipy.optimize as op

def f_cog_generate(coeffs):
    theta = sp.symbols('theta')
    return np.sum([coeff[0]*sp.sin(coeff[1]*theta+coeff[2]) for coeff in coeffs])

def my_fft(Ts, y:np.ndarray):
   sp = np.fft.fft(y)
   freq = np.fft.fftfreq(y.shape[-1])
   amplitude = (np.abs(sp)/len(y)*2)[:int(len(freq)/2)]
   freq_Hz = (freq/Ts)[:int(len(freq)/2)]
   return freq_Hz, amplitude

# end def
if __name__ == "__main__":
    pass
    # %% [markdown]
    # 电流波动假数据
    theta = sp.symbols('theta')
    f_cog_expr = f_cog_generate([
        [0.05,1,0],
        [0.1,3,0],
        [0.03,5,0],
        [0.01,10,0],
    ])
    Ts = 1e-4
    t = np.arange(0, 2, Ts, dtype=float) 
    theta_data = 2*np.pi*t
    f_cog_data = [float(f_cog_expr.evalf(subs={theta: data})) for data in theta_data] + np.random.uniform(-0.0, 0.0, len(theta_data))
    # theta_data = [float(data) for data in theta_data]
    # f_cog_data = [float(data) for data in f_cog_data]
    plt.scatter(t, f_cog_data, marker='.')
    plt.show()
    
    # %% [markdown]
    # 读取并处理真实电流波动数据  D:\BaiduSyncdisk\PersonalDataBase\1_WorkDatabase\2_CXSW\1_Project\1_StepperMotorControl\4_Data
    with open('../../../4_Data/spd_data.txt', 'r') as f:
        line = f.readline().splitlines()
        
        i=0
        data_list = []
        while line:
            i+=1
            data = float(*line)
            data_list.append(data)
            # print ("{}: {} | {}".format(i, data, type(data)))
            line = f.readline().splitlines()
        print("size:{}".format(len(data_list)))
    
    data_list = np.array(data_list)
    data_list_mean = np.mean(data_list)
    data_list_denoising = data_list[abs(data_list-data_list_mean)<100] - data_list_mean
    # data_list_denoising = data_list[abs(data_list-data_list_mean)<50]
    Ts = 1e-4
    x_data = np.arange(len(data_list_denoising))*Ts
    plt.plot(x_data, data_list_denoising)
    plt.show()
    
    # %% [markdown]
    # 读取xlxs数据
    # import xlrd
    # wd = xlrd.open_workbook('D:/BaiduSyncdisk/PersonalDataBase/1_WorkDatabase/2_CXSW/1_Project/1_StepperMotorControl/4_Data/s1.xlsx')
    import pandas as pd
    # df = pd.read_excel('../../../4_Data/s1.xlsx')
    df = pd.read_excel('../../../4_Data/a.xlsx')
    data_start = 1000
    data_end = 7000
    Ts = 1e-4
    pos_data = df.values[:,0][data_start:data_end]
    iq_data = df.values[:,1][data_start:data_end]
    t_data = np.arange(len(pos_data)) * Ts
    
    iq_data = iq_data - np.mean(iq_data)
    plt.plot(pos_data, iq_data)
    plt.show()
    
    T_pos = (max(pos_data) - min(pos_data)) / len(pos_data)
    
    # %% [markdown]
    # 扭矩波动fft分析
    Ts = 1e-4
    # y_data = iq_data
    y_data = f_cog_data
    freq, amp = my_fft(Ts, y_data) 
    nfreq = 2*np.pi*freq
    
    amp_threshold = 0.01
    coeff_reconstruct = np.array([amp[amp>amp_threshold],
                            nfreq[amp>amp_threshold],
                                np.zeros(len(amp[amp>amp_threshold]))]).T
    print("coeff_reconstruct:{}".format(coeff_reconstruct))
    
    plt.plot(nfreq, amp)
    plt.xlim([0,1000])
    plt.grid(visible=True)
    ax = plt.gca()
    # ax.yaxis.set_major_locator(plt.MultipleLocator(0.01))
    plt.show()  
    
    
    # %% [markdown]
    # 扭矩波动重构： 根据fft结果，取幅值占主导的频率成分重构函数
    theta = sp.symbols('theta')
    Ts = 1e-4
    theta_data = 2*np.pi*t
    t = np.arange(0, len(y_data)) * Ts
    theta_data = 2*np.pi*t
    # theta_data = theta_data[:int(len(x_data)*1)]
    f_cog_expr_conrestruct = f_cog_generate(coeff_reconstruct)
    
    from time import time
    start = time()
    # f_cog_data_conrestruct = [float(f_cog_expr_conrestruct.subs(theta, data)) for data in theta_data]
    f_cog_data_conrestruct = [float(f_cog_expr_conrestruct.evalf(subs={theta: data})) for data in theta_data]
    end = time()
    print("cost:{}s".format(end-start))
    
    plt.plot(theta_data, f_cog_data_conrestruct, marker='.')
    # plt.xlim([0,1000])
    plt.show()
    
    # %% [markdown]
    # 扭矩波动补偿（函数法）：把转子位置代入拟合出的扭矩波动补偿表达式，计算出需要补偿的电流值（iq: A）
    